from langgraph.graph import StateGraph,START,END
from typing_extensions import TypedDict,Annotated,List
from langchain_openai import ChatOpenAI
import operator
from langchain_core.messages import HumanMessage,SystemMessage,AIMessage

llm = ChatOpenAI(
    model='deepseek-chat',
    base_url='https://api.deepseek.com/v1',
    api_key='sk-a51c735868b94a03be98ed0e3b233842'
)
class State(TypedDict):
  messages:Annotated[List[str],operator.add]
bulier = StateGraph(State)
def chat_node(State):
  print(State)
  print('-'*20)
  messages = State['messages']
  llm_answer = llm.invoke(messages).content
  return {'messages':[llm_answer]}

def convert_node(State):
  print(State)
  print('-'*20)
  prompt ="您是一位数据提取专家，负责从文本中检索关键信息。请为所提供的文本提取相关信息，并以 JSON 格式输出。概述所提取的关键数据点。"
  messages=[SystemMessage(content=prompt),HumanMessage(content=State['messages'][-1])]
  res = llm.invoke(messages).content
  return {'messages':[res]}
bulier.add_node('对话模型',chat_node)
bulier.add_node('数据提取模型',convert_node)
bulier.add_edge(START,'对话模型')
bulier.add_edge('对话模型','数据提取模型')
bulier.add_edge('数据提取模型',END)
graph = bulier.compile()
query='帮我生成一首李白风格的诗，主题为‘月夜编程'
mes= {'messages':[HumanMessage(content=query)]}
res=graph.invoke(mes)
print(res['messages'][-1])



              

